'''
给出模块中的不同模块的权重值，并生成对应的md文件
注意这里简单的输出固定格式的md格式文件
因此直接使用简单的元素创建
'''
import os
import sys
import numpy as np
import matplotlib
from matplotlib import pyplot as plt
import torch
import math
import json

from torch.serialization import save

def CreateFolder(folderpath):
    if not os.path.exists(folderpath):
        os.makedirs(folderpath)

def CreateImageMeta(img_path,height,width):
    '''
    创建一个img 元素对象
    args:
        img_path: 图片的路径
        height： 高度
        width：宽度
    '''
    img_str="<img src='{}' width='{}px' height='{}px'>".format(img_path,height,width)
    return img_str

def CreateMatrix_double(weight_matrix):
    '''
    输出权重矩阵，对应的数学latex字符
    '''

    matrix_content=""
    h,w=weight_matrix.shape # 获取矩阵形状 行，列
    for i in range(h):
        for j in range(w):
            matrix_content="{} & {}".format(matrix_content,'%.5f'%weight_matrix[i,j])
        matrix_content=r"{} \\ ".format(matrix_content)
    matrix_str=r'$$\begin{bmatrix} '+matrix_content+ r'\end{bmatrix}$$'
    return matrix_str

def CreateMatrix_single(weight_matrix):
    '''
    输出权重矩阵，对应的数学latex字符串
    '''
    matrix_content=""
    h,w=weight_matrix.shape # 获取矩阵形状 行，列
    for i in range(h):
        for j in range(w):
            matrix_content="{} & {}".format(matrix_content,'%.5f'%weight_matrix[i,j])
    matrix_str=r'$$\begin{bmatrix} '+matrix_content+ r'\end{bmatrix}$$'
    return matrix_str

def SaveNumpy_Weight(root_dir,Convlayer_dict):
    '''
    保存对应的权重，并将权重文件存放在root_dir
    输出：
        1、模块的权重文件
            a.卷积直接输出卷积核的权重文件
            b.正则化，直接按照名字输出保存文件
            c.激活函数没有权重，所以只存在一个文件夹名称
    '''
    for conv_name in Convlayer_dict:
        for conv_module_name in Convlayer_dict[conv_name]:
            savepath=os.path.join(root_dir,conv_name,conv_module_name)
            CreateFolder(savepath)
            if conv_module_name=="conv": # 卷积核
                # 直接保存结果
                savepath=os.path.join(savepath,'weight.npy')
                np.save(savepath,Convlayer_dict[conv_name][conv_module_name].weight.detach().cpu().numpy())               
            elif conv_module_name=="batchnorm": # 正则化
                # 保存正则化中的各种参数
                bias_path=os.path.join(savepath,"bias.npy")
                gamma_path=os.path.join(savepath,"gamma.npy")
                mean_path=os.path.join(savepath,"mean.npy")
                var_path=os.path.join(savepath,"var.npy")
                np.save(bias_path,Convlayer_dict[conv_name][conv_module_name].bias.detach().cpu().numpy()) 
                np.save(gamma_path,Convlayer_dict[conv_name][conv_module_name].weight.detach().cpu().numpy()) 
                np.save(mean_path,Convlayer_dict[conv_name][conv_module_name].running_mean.detach().cpu().numpy()) 
                np.save(var_path,Convlayer_dict[conv_name][conv_module_name].running_var.detach().cpu().numpy()) 
            elif conv_module_name=="ReLU": # 激活函数
                pass

def ConvTable(Conv_module_dict,CreateMatrix):
    '''
    根据当前模块的结果，创建对应的表格
    args:
        Conv_module:当前卷积模块情况
    1、保存权重对应的渲染图片，并将结果保存成单张图片
        a.渲染单张图片，并保存至特定文件夹中
        b.
        c.保存贡献度对比图
        d.
    '''
    # 创建表头
    conv_b,conv_c,conv_h,conv_w=Conv_module_dict['conv'].weight.shape
    # 输出波段，输入波段，
    bias_batch=Conv_module_dict['batchnorm'].bias.detach().cpu().numpy()
    mean_batch=Conv_module_dict['batchnorm'].running_mean.detach().cpu().numpy()
    var_batch=Conv_module_dict['batchnorm'].running_var.detach().cpu().numpy()
    gamma_batch=Conv_module_dict['batchnorm'].weight.detach().cpu().numpy()
    table_head="|channel_index|"
    table_alignment="|:-|"
    for i in range(conv_c):
        table_head="{}channel_{}|".format(table_head,str(i)) # 确定表头
        table_alignment="{}:---|".format(table_alignment)
    table_head="{}gamma|bias|mean|var|\n".format(table_head)
    table_alignment="{}:--|:--|:--|:--|\n".format(table_alignment)
    # 确定表格的内容
    table_rows=[]
    # 确定表格的行列情况
    for i in range(conv_b):
        table_row="|{}|".format(i) # 单行
        for j in range(conv_c):
            k_weight=Conv_module_dict['conv'].weight[i,j,:,:].detach().cpu().numpy()
            conv_k_weight_str=CreateMatrix(k_weight) # 创建对应的权重的结果
            table_row="{}{}|".format(table_row,conv_k_weight_str)
        # 确定对应行的batchnorm结构
        table_row='{}{}|'.format(table_row,gamma_batch[i])
        table_row='{}{}|'.format(table_row,bias_batch[i])
        table_row='{}{}|'.format(table_row,mean_batch[i])
        table_row='{}{}|'.format(table_row,var_batch[i])
        table_row='{}\n'.format(table_row)
        table_rows.append(table_row)
    # 构建完整矩阵结构
    table_str=""
    table_str="{}{}".format(table_head,table_alignment)
    for row in table_rows:
        table_str='{}{}'.format(table_str,row)
    return table_str # 确定输出表格的字符文字

def ConvTable_double(Conv_module_dict):
    '''双行矩阵'''
    return ConvTable(Conv_module_dict,CreateMatrix_double)


def ConvTable_single(Conv_module_dict):
    '''单行矩阵'''
    return ConvTable(Conv_module_dict,CreateMatrix_single)

def CreateMarkDown(rootdir,module_dict):
    '''创建对应的markdown文件，用于确定markdown'''
    markdown_str=""
    for conv_key in module_dict:
        markdown_str="{}\n### double shape matrix  {}".format(markdown_str,conv_key) # 
        markdown_str="{}\n{}".format(markdown_str,ConvTable_double(module_dict[conv_key]))
    for conv_key in module_dict:
        markdown_str="{}\n### single shape matrix  {}".format(markdown_str,conv_key) # 
        markdown_str="{}\n{}".format(markdown_str,ConvTable_single(module_dict[conv_key]))

    # 保存对应的图形
    with open(os.path.join(rootdir,"weight.md"),'w',encoding='utf-8') as fp:
        fp.write(markdown_str)
    
def RanderTable():
    '''
    渲染表格以及确定该表格的情况
    '''
    pass

if __name__=="__main__":
    weight_matrix=np.ones((3,3))
    print(CreateMatrix_double(weight_matrix))